# -*- coding: utf-8 -*-
# @Author   : zxc19
# @Time     : 2023/7/6 21:10
# @FileName : ssh_tunnel.py
# @Desc     : 请输入功能模块描述
from typing import Optional

from loguru import logger
from pydantic import Field, BaseModel
from sshtunnel import SSHTunnelForwarder


class SSHTunnelObj(BaseModel):
    """
    SSH隧道对象模型
    """
    ssh_host: str = Field('127.0.0.1', description='堡垒机ip地址或主机名')
    ssh_port: int = Field(22, description='堡垒机连接mysql服务器的端口号，一般都是22，必须是数字')
    ssh_username: str = Field('root', description='堡垒机用户名')
    ssh_password: str = Field('123456', description='堡垒机用户密码')
    remote_bind_host: str = Field(..., description='远程主机地址')
    remote_bind_port: int = Field(..., description='远程主机端口')
    bind_local_address: bool = Field(False, description='是否指定本地地址和端口')
    local_bind_host: Optional[str] = Field(None, description='使用堡垒机的哪张网卡进行连接')
    local_bind_port: Optional[int] = Field(None, description='使用堡垒机的哪个端口进行连接')


class SSHTunnel(object):
    """
    SSH隧道
    """

    def __init__(self, ssh_tunnel: SSHTunnelObj):
        """
        初始化并开启ssh隧道
        :param ssh_tunnel:
        """
        self._ssh_tunnel = ssh_tunnel
        logger.debug(f'正在初始化SSH隧道：{self._ssh_tunnel.dict()}')
        self._server = SSHTunnelForwarder(
            ssh_address_or_host=(self._ssh_tunnel.ssh_host, self._ssh_tunnel.ssh_port),
            ssh_username=self._ssh_tunnel.ssh_username,
            ssh_password=self._ssh_tunnel.ssh_password,
            remote_bind_address=(self._ssh_tunnel.remote_bind_host, self._ssh_tunnel.remote_bind_port),
        )
        self._server.start()
        logger.debug(f'SSH隧道初始化完成：{self._ssh_tunnel.dict()}')

    @property
    def server(self):
        """
        获取隧道
        :return:
        """
        return self._server

    def close(self):
        """
        关闭隧道
        :return:
        """
        logger.debug(f'正在关闭SSH隧道：{self._ssh_tunnel.dict()}')
        self._server.close()
        logger.debug(f'SSH隧道已关闭：{self._ssh_tunnel.dict()}')

